# -*- coding: utf-8 -*-
"""
修正：人均可支配收入是“季度累计”，先还原为“单季收入”，再与就业率做相关性分析
输出：
  - corr_income_levels.csv   # 单季收入 vs 就业率(季度均值)
  - corr_income_yoy.csv      # 单季收入YoY vs 就业率YoY
"""

import os
import numpy as np
import pandas as pd
from scipy.stats import pearsonr, spearmanr
from statsmodels.stats.multitest import multipletests

# ========= 改这里的路径 =========
PATH_EMP = "/Users/linshangjin/25CCM/NKU-C/dataset_csv/data_processed.xlsx"   # 就业率预测（多Sheet，date, employment_rate）
PATH_INC = "/Users/linshangjin/25CCM/NKU-C/dataset_csv/人均收入_季度数据.csv"       # 季度GDP与产业结构
OUT_DIR  = "/Users/linshangjin/25CCM/NKU-C/t2/人均收入"                              # 输出目录
# =================================

os.makedirs(OUT_DIR, exist_ok=True)

# ---------- 工具 ----------
def read_employment_xlsx(path):
    xl = pd.ExcelFile(path)
    parts = []
    for sh in xl.sheet_names:
        t = xl.parse(sh)
        if {"date","employment_rate"}.issubset(t.columns):
            t = t[["date","employment_rate"]].copy()
            t["date"] = pd.to_datetime(t["date"])
            t["employment_rate"] = pd.to_numeric(t["employment_rate"], errors="coerce")
            parts.append(t)
    df = pd.concat(parts, ignore_index=True).dropna(subset=["date"]).sort_values("date")
    return df.set_index("date").sort_index()

def read_income_csv(path):
    # 允许日期列为 'date' 或 'Unnamed: 0'；其余列都当作数值（可多列）
    inc = pd.read_csv(path, encoding="utf-8")
    if "date" in inc.columns:
        pass
    elif "Unnamed: 0" in inc.columns:
        inc = inc.rename(columns={"Unnamed: 0":"date"})
    else:
        raise ValueError("收入文件缺少日期列（date 或 Unnamed: 0）")
    inc["date"] = pd.to_datetime(inc["date"])
    for c in inc.columns:
        if c != "date":
            inc[c] = pd.to_numeric(inc[c], errors="coerce")
    return inc.set_index("date").sort_index()

def within_year_quarter_flow_from_cum(s):
    """
    输入：季度“累计”序列（通常每年Q1重置）
    输出：同年内差分得到“单季收入”，Q1 用当季累计值
    """
    sq = s.resample("Q").last()            # 取季末（确保是季度频）
    grp = sq.groupby(sq.index.year)
    flow = grp.apply(lambda x: x.diff().fillna(x))  # Q2-4 = diff, Q1 = 本季累计
    flow.index = sq.index                  # restore DatetimeIndex
    return flow

def corr_with_p(a, b, method="pearson"):
    idx = a.dropna().index.intersection(b.dropna().index)
    a1, b1 = a.loc[idx], b.loc[idx]
    if len(a1) < 8:
        return np.nan, np.nan, len(a1)
    if method == "pearson":
        r, p = pearsonr(a1, b1)
    else:
        r, p = spearmanr(a1, b1)
    return float(r), float(p), len(a1)

# ---------- 读数据 ----------
emp = read_employment_xlsx(PATH_EMP)   # 日/半月就业率
y_q  = emp["employment_rate"].resample("Q").mean().rename("就业率_季度均值").dropna()

inc_raw = read_income_csv(PATH_INC)    # 季度“累计”收入（可能不止一列）
num_cols = [c for c in inc_raw.columns if c != "date"]  # 所有数值列

# ---------- 从“累计”还原“单季收入” ----------
flows = {}
for col in num_cols:
    flows[col + "_单季收入"] = within_year_quarter_flow_from_cum(inc_raw[col])

inc_flow = pd.DataFrame(flows)

# ---------- 对齐 ----------
joint_lvl = inc_flow.join(y_q, how="inner").dropna()

# ---------- 1) 单季“水平”相关 ----------
rows_lvl = []
for col in inc_flow.columns:
    rP, pP, n = corr_with_p(joint_lvl[col], joint_lvl["就业率_季度均值"], method="pearson")
    rS, pS, _ = corr_with_p(joint_lvl[col], joint_lvl["就业率_季度均值"], method="spearman")
    rows_lvl.append({"variable": col, "pearson": rP, "pearson_p": pP,
                     "spearman": rS, "spearman_p": pS, "n_obs": n})
tab_lvl = pd.DataFrame(rows_lvl).sort_values("pearson", ascending=False)
mask = tab_lvl["pearson_p"].notna()
if mask.any():
    _, qvals, *_ = multipletests(tab_lvl.loc[mask,"pearson_p"], method="fdr_bh")
    tab_lvl.loc[mask,"pearson_q"] = qvals
lvl_path = os.path.join(OUT_DIR, "corr_income_levels.csv")
tab_lvl.to_csv(lvl_path, index=False)

# ---------- 2) YoY（4期差分）相关 ----------
y_yoy = y_q.diff(4).rename("就业率_YoY")
rows_yoy = []
for col in inc_flow.columns:
    x_yoy = inc_flow[col].diff(4)   # 单季收入同比（不用log，避免零/负值出错）
    # 对齐
    idx = x_yoy.dropna().index.intersection(y_yoy.dropna().index)
    if len(idx) < 8:
        rows_yoy.append({"variable": col, "pearson_yoy": np.nan, "pearson_p": np.nan,
                         "spearman_yoy": np.nan, "spearman_p": np.nan, "n_obs": len(idx)})
        continue
    rP, pP = pearsonr(x_yoy.loc[idx], y_yoy.loc[idx])
    rS, pS = spearmanr(x_yoy.loc[idx], y_yoy.loc[idx])
    rows_yoy.append({"variable": col, "pearson_yoy": float(rP), "pearson_p": float(pP),
                     "spearman_yoy": float(rS), "spearman_p": float(pS), "n_obs": int(len(idx))})
tab_yoy = pd.DataFrame(rows_yoy).sort_values("pearson_yoy", ascending=False)
mask2 = tab_yoy["pearson_p"].notna()
if mask2.any():
    _, qvals2, *_ = multipletests(tab_yoy.loc[mask2,"pearson_p"], method="fdr_bh")
    tab_yoy.loc[mask2,"pearson_q"] = qvals2
yoy_path = os.path.join(OUT_DIR, "corr_income_yoy.csv")
tab_yoy.to_csv(yoy_path, index=False)

print("✅ 已输出：")
print(" -", lvl_path, "(单季收入水平相关)")
print(" -", yoy_path, "(单季收入同比相关)")
